package com.cloud.zuul.core;

import com.cloud.zuul.context.VersionHolder;
import com.netflix.loadbalancer.RoundRobinRule;
import com.netflix.loadbalancer.Server;
import com.netflix.zuul.context.RequestContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.netflix.zuul.filters.support.FilterConstants;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
@Scope(ConfigurableBeanFactory.SCOPE_PROTOTYPE)
public class RibbonServerSelector extends RoundRobinRule {

	@Autowired
	private VersionHolder versionHolder;

	@Autowired
	private DiscoveryClient discoveryClient;

	private Map<String, Integer> selectIndexMap = new HashMap<>();

	@Override
	public Server choose(Object key) {
		RequestContext context = RequestContext.getCurrentContext();
		String serviceId = (String) context.get(FilterConstants.SERVICE_ID_KEY);

		String version = versionHolder.get();
		if (version != null) {
			List<ServiceInstance> matchInstanceList = new ArrayList<>();
			List<ServiceInstance> instanceList = discoveryClient.getInstances(serviceId);
			for (ServiceInstance instance : instanceList) {
				String instanceVersion = instance.getMetadata().get("version");
				if (StringUtils.isEmpty(instanceVersion) || instanceVersion.equals("undefined") || !instanceVersion.equals(version)) {
					continue;
				}
				matchInstanceList.add(instance);
			}
			if (matchInstanceList.size() > 0) {
				//选择服务器
				ServiceInstance instance = selectInstance(serviceId, matchInstanceList);
				List<Server> servers = getLoadBalancer().getAllServers();
				for (Server server : servers) {
					if (server.getHost().equals(instance.getHost()) && server.getPort() == instance.getPort()) {
						return server;
					}
				}
			}
		}
		return super.choose(key);
	}

	private synchronized ServiceInstance selectInstance(String serviceId, List<ServiceInstance> matchInstanceList) {
		Integer selectIndex = selectIndexMap.get(serviceId);
		if (selectIndex == null || selectIndex == Short.MAX_VALUE) {
			selectIndex = 0;
		}
		selectIndexMap.put(serviceId, ++selectIndex);
		return matchInstanceList.get(selectIndex % matchInstanceList.size());
	}

}
